# 导入增强版知识图谱提取器
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'datagraph'))

try:
    from kg_extractor import EnhancedKGTripletExtractor
    KGTripletExtractor = EnhancedKGTripletExtractor
except ImportError:
    # 如果导入失败，保留原有的简单版本
    import jieba
    import jieba.posseg as pseg
    import re
    from collections import defaultdict
    
    class KGTripletExtractor:
        def __init__(self):
            self.load_resources()
            
        def load_resources(self):
            """加载必要的语言资源"""
            jieba.add_word('清华大学', freq=1000, tag='nt')
            jieba.add_word('北京大学', freq=1000, tag='nt')
            jieba.add_word('中国科学院', freq=1000, tag='nt')
            
            self.relation_words = {
                '出生', '出生于', '发明', '发现', '创作', '撰写', '写', '出版',
                '担任', '就职于', '工作于', '毕业于', '就读于', '来自', '位于'
            }
            
            self.stop_words = {'的', '了', '和', '与', '及', '在', '是', '有', '我', '你', '他', '她', '它'}
    
    def extract_triplets(self, text):
        """
        从中文文本中提取三元组
        :param text: 输入文本
        :return: 三元组列表 [(subject, relation, object)]
        """
        sentences = self.split_sentences(text)
        triplets = []
        
        for sentence in sentences:
            # 使用jieba进行分词和词性标注
            words = pseg.cut(sentence)
            words_list = list(words)
            
            # 提取可能的三元组
            sentence_triplets = self.extract_from_sentence(words_list)
            triplets.extend(sentence_triplets)
            
        return triplets
    
    def split_sentences(self, text):
        """将文本分割成句子"""
        # 简单的中文分句规则
        sentences = re.split(r'[。！？!?;；]', text)
        return [s.strip() for s in sentences if len(s.strip()) > 0]
    
    def extract_from_sentence(self, words_list):
        """
        从单个句子中提取三元组
        :param words_list: 分词和词性标注结果列表
        :return: 三元组列表
        """
        triplets = []
        
        # 寻找可能的关系词
        relation_positions = []
        for i, (word, flag) in enumerate(words_list):
            if word in self.relation_words:
                relation_positions.append(i)
        
        # 对每个关系词，尝试提取其前后的主体和客体
        for rel_pos in relation_positions:
            rel_word = words_list[rel_pos][0]
            
            # 提取关系词前的主体（主语）
            subject = self.extract_subject(words_list, rel_pos)
            
            # 提取关系词后的客体（宾语）
            obj = self.extract_object(words_list, rel_pos)
            
            if subject and obj:
                triplets.append((subject, rel_word, obj))
        
        return triplets
    
    def extract_subject(self, words_list, rel_pos):
        """提取关系词前的主体"""
        subject_words = []
        
        # 从关系词向前遍历，收集名词性成分作为主体
        for i in range(rel_pos-1, -1, -1):
            word, flag = words_list[i]
            
            # 如果是停用词，停止收集
            if word in self.stop_words:
                break
                
            # 收集名词性成分（名词、人名、地名、机构名等）
            if flag.startswith('n') or flag == 'nr' or flag == 'ns' or flag == 'nt':
                subject_words.insert(0, word)  # 逆序插入保持原顺序
            else:
                break
                
        return ''.join(subject_words) if subject_words else None
    
    def extract_object(self, words_list, rel_pos):
        """提取关系词后的客体"""
        object_words = []
        
        # 从关系词向后遍历，收集名词性成分作为客体
        for i in range(rel_pos+1, len(words_list)):
            word, flag = words_list[i]
            
            # 如果是停用词，停止收集
            if word in self.stop_words:
                break
                
            # 收集名词性成分
            if flag.startswith('n') or flag == 'nr' or flag == 'ns' or flag == 'nt':
                object_words.append(word)
            else:
                break
                
        return ''.join(object_words) if object_words else None

# 测试代码
if __name__ == "__main__":
    extractor = KGTripletExtractor()
    
    # 扩展的测试用例（包含中英文混合）
    test_sentences = [
        # 原有测试用例
        "马云出生于浙江省杭州市。",
        "曹雪芹创作了红楼梦。",
        "爱因斯坦发现了相对论。",
        "清华大学位于北京市。",
        "莫言毕业于北京大学并获得了诺贝尔文学奖。",
        
        # 新增英文和混合测试用例
        "苹果公司发布了iPhone 15，这款手机搭载了A17芯片。",
        "OpenAI开发了ChatGPT，这是一个基于GPT-4的AI助手。",
        "Elon Musk创立了Tesla和SpaceX两家公司。",
        "GitHub被Microsoft收购，价格为75亿美元。",
        "Python是由Guido van Rossum在1991年创建的编程语言。",
        "Google的搜索引擎使用了PageRank算法。",
        "华为在2019年发布了HarmonyOS操作系统。",
        "阿里云提供云计算服务，支持Docker容器技术。",
        "腾讯QQ有超过8亿用户。",
        "清华大学与MIT建立了合作关系。"
    ]
    
    print("知识图谱三元组提取测试（增强版）")
    print("=" * 60)
    
    total_triplets = 0
    for i, sentence in enumerate(test_sentences, 1):
        print(f"{i}. 输入: {sentence}")
        
        # 使用增强版提取器
        if hasattr(extractor, 'extract_triplets'):
            triplets = extractor.extract_triplets(sentence, normalize_relations=False)
        else:
            # 回退到简单版本
            triplets = extractor.extract_triplets(sentence)
        
        if triplets:
            for j, (s, r, o) in enumerate(triplets, 1):
                print(f"   三元组 {j}: ({s}, {r}, {o})")
                total_triplets += 1
        else:
            print("   未提取到三元组")
        print()
    
    print(f"总计提取到 {total_triplets} 个三元组")
    
    # 如果使用增强版提取器，展示更多功能
    if hasattr(extractor, 'extract_enhanced_entities'):
        print("\n" + "=" * 60)
        print("实体识别测试（增强版）")
        print("=" * 60)
        
        test_text = "苹果公司的CEO Tim Cook在2023年宣布，iPhone 15将支持USB-C接口，这标志着苹果向统一标准迈出了重要一步。该产品预计将在apple.com上发售。"
        print(f"测试文本: {test_text}")
        print("\n识别到的实体:")
        
        entities = extractor.extract_enhanced_entities(test_text)
        for entity in entities:
            if isinstance(entity, dict):
                print(f"  - {entity['text']} (类型: {entity['type']}, 来源: {entity['source']})")
            else:
                print(f"  - {entity}")
                
        print(f"\n该文本的三元组提取结果:")
        triplets = extractor.extract_triplets(test_text, normalize_relations=False)
        for j, (s, r, o) in enumerate(triplets, 1):
            print(f"   三元组 {j}: ({s}, {r}, {o})")